#include <iostream>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <future>
#include <memory>
#include <vector>
#include <stdexcept>

/*
 * Thread pool class.
 */

class ThreadPool {
public:
    explicit ThreadPool(size_t threads) : m_stop(false) {
        if(threads == 0) {
            threads = 1;
        }
        
        // Start to run threads.
        for(size_t i = 0; i < threads; ++i) {
            m_workers.emplace_back(
                [this] {
                    for(;;) {
                        std::function<void()> task;
                        
                        // Get task object.
                        {
                            std::unique_lock<std::mutex> lock(this->m_mutex);
                            this->m_condition.wait(lock, [this]{ return this->m_stop || !this->m_tasks.empty(); });     
                            if(this->m_stop && this->m_tasks.empty()) {
                                return;
                            }          
                            task = std::move(this->m_tasks.front());
                            this->m_tasks.pop();
                        }
                        
                        // Execute task.
                        task();
                    }
                }
            );
        }
    }
     ~ThreadPool() {
        // Try to set stop flag.
        {
            std::unique_lock<std::mutex> lock(m_mutex);
            m_stop = true;
        }

        // Notify all threads to stop.
        m_condition.notify_all();
        for(std::thread &worker: m_workers) {
            worker.join();
        }
    }

    // push task function.
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args) -> std::future<typename std::result_of<F(Args...)>::type> {
        using return_type = typename std::result_of<F(Args...)>::type;
        auto task = std::make_shared<std::packaged_task<return_type()> >(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...)
        );
            
        std::future<return_type> res = task->get_future();
        {
            std::unique_lock<std::mutex> lock(m_mutex);
            if(m_stop) {
                throw std::runtime_error("enqueue on m_stopped ThreadPool");
            }              
            m_tasks.push([task](){ (*task)(); });
        }
        m_condition.notify_one();
        return res;
    }

    // size just get task size.
    size_t size() const {
        return m_workers.size();
    }

private:
    std::vector<std::thread> m_workers;
    std::queue<std::function<void()>> m_tasks;
    std::mutex m_mutex;
    std::condition_variable m_condition;
    bool m_stop;
};

/*
// 示例用法
int main() {
    try {
        // 创建线程池，使用4个工作线程
        ThreadPool pool(4);
        std::cout << "Thread pool created with " << pool.size() << " threads" << std::endl;
        
        // 提交多个任务
        std::vector<std::future<int>> results;
        
        for(int i = 0; i < 8; ++i) {
            results.emplace_back(
                pool.enqueue([i] {
                    std::cout << "Task " << i << " started in thread " 
                              << std::this_thread::get_id() << std::endl;
                    std::this_thread::sleep_for(std::chrono::seconds(1));
                    std::cout << "Task " << i << " finished" << std::endl;
                    return i * i;
                })
            );
        }
        
        // 获取结果
        for(size_t i = 0; i < results.size(); ++i) {
            std::cout << "Result of task " << i << ": " << results[i].get() << std::endl;
        }
        
    } catch(const std::exception& e) {
        std::cerr << "Exception: " << e.what() << std::endl;
    }
    
    return 0;
}*/